# TlpDs测试类
import unittest
from torch.utils.data import Dataset, DataLoader
from apps.tlp.tlp_ds import TlpDs
from apps.tlp.tlp_config import TlpConfig as TG

class TTlpDs(unittest.TestCase):
    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_seconds_between
    def test_seconds_between(self):
        t1 = '2025-02-12 10:16:00'
        t2 = '2025-02-12 10:17:32'
        td = TlpDs.seconds_between(t1, t2)
        print(f'td: {td};')

    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_load_result_csv
    def test_load_result_csv(self):
        csv_fn = 'work/tlp/datas/warnings1-8.csv'
        rec_dict = TlpDs.load_result_csv(csv_fn=csv_fn)
        num = 0
        for k, v in rec_dict.items():
            print(f'### {k}: {v};')
            num += 1
            if num > 20:
                break

    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_extract_frames
    def test_extract_frames(self):
        '''
        '''
        raw_datas = [
            '188.0', '2.0', '3.0', '4.0', '5.0',
            '6.0', '7.0', '8.0', '9.0', '10.0',
            '11.0', '12.0', '13.0', '14.0', '15.0',
            '16.0', '17.0', '18.0', '19.0', '20.0'
        ]
        TlpDs.FRAME_SIZE = 3
        TlpDs.fRAME_STEP = 2
        csv_fn = 'work/tlp/datas/warnings1-8.csv'
        rec_dict = TlpDs.load_result_csv(csv_fn=csv_fn)
        Xs, ys = TlpDs.extract_frames(raw_datas=raw_datas, rec=rec_dict[str(3)])
        print(f'Xs: {Xs.shape}; ys: {ys.shape};')
        print(f'{Xs};')

    # python -m unittest -v -k test_classify_frame_ uts.apps.tlp.t_tlp_ds.TTlpDs
    def test_classify_frame_001(self):
        # 没有雷电
        start_idx, end_idx = 100, 200
        tl_start_idx, tl_end_idx = -1, -1
        cid = TlpDs.classify_frame(start_idx, end_idx, tl_start_idx, tl_end_idx)
        print(f'cid={cid};')
        self.assertTrue(cid==0, msg='不能分类正常情况（无雷电）')
    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_classify_frame_002
    def test_classify_frame_002(self):
        start_idx, end_idx = 100, 200
        tl_start_idx, tl_end_idx = TG.WARN_YELLOW + end_idx + 100, TG.WARN_YELLOW + start_idx + 1000
        cid = TlpDs.classify_frame(start_idx, end_idx, tl_start_idx, tl_end_idx)
        print(f'cid={cid};')
        self.assertTrue(cid==0, msg='不能分类正常情况（30分钟内无雷电）')
    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_classify_frame_003
    def test_classify_frame_003(self):
        start_idx, end_idx = 100, 200
        tl_start_idx, tl_end_idx = TG.WARN_RED + end_idx - 1, TG.WARN_YELLOW + end_idx + 1000
        cid = TlpDs.classify_frame(start_idx, end_idx, tl_start_idx, tl_end_idx)
        print(f'cid={cid};')
        self.assertTrue(cid==3, msg='不能分类红色预警')
    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_classify_frame_004
    def test_classify_frame_004(self):
        # 没有雷电
        start_idx, end_idx = 100, 200
        tl_start_idx, tl_end_idx = TG.WARN_BROWN + end_idx - 1, TG.WARN_YELLOW + end_idx + 1000
        cid = TlpDs.classify_frame(start_idx, end_idx, tl_start_idx, tl_end_idx)
        print(f'cid={cid};')
        self.assertTrue(cid==2, msg='不能分类橙色预警')
    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_classify_frame_005
    def test_classify_frame_005(self):
        start_idx, end_idx = 100, 200
        tl_start_idx, tl_end_idx = TG.WARN_YELLOW + end_idx - 1, TG.WARN_YELLOW + end_idx + 1000
        cid = TlpDs.classify_frame(start_idx, end_idx, tl_start_idx, tl_end_idx)
        print(f'cid={cid};')
        self.assertTrue(cid==1, msg='不能分类黄色预警')
    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_classify_frame_006
    def test_classify_frame_006(self):
        base_idx = 1000
        start_idx, end_idx = base_idx + 10, base_idx + 200
        tl_start_idx, tl_end_idx = base_idx, base_idx + 1000
        cid = TlpDs.classify_frame(start_idx, end_idx, tl_start_idx, tl_end_idx)
        print(f'cid={cid};')
        self.assertTrue(cid==4, msg='不能分类开始时间在雷电期间情况')
    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_classify_frame_007
    def test_classify_frame_007(self):
        # 没有雷电
        base_idx = 1000
        start_idx, end_idx = base_idx + 10, base_idx + 200
        tl_start_idx, tl_end_idx = base_idx, base_idx + 1000
        cid = TlpDs.classify_frame(start_idx, end_idx, tl_start_idx, tl_end_idx)
        print(f'cid={cid};')
        self.assertTrue(cid==4, msg='不能分类结束时间在雷电期间情况')

    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_process_raw_datas
    def test_process_raw_datas(self):
        txt_fn = 'work/tlp/datas/electric.txt'
        csv_fn = 'work/tlp/datas/warnings1-8.csv'
        rec_dict = TlpDs.load_result_csv(csv_fn=csv_fn)
        TlpDs.process_raw_datas(txt_fn=txt_fn, rec_dict=rec_dict)

    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_cls_init
    def test_cls_init(self):
        ds = TlpDs()
        print(f'数据集大小：{len(ds)}; ????????????')

    # python -m unittest -v uts.apps.tlp.t_tlp_ds.TTlpDs.test_dataloader
    def test_dataloader(self):
        ds = TlpDs()
        Xi, yi = ds[25321]
        print(f'Xi: {Xi.shape}; yi: {yi.shape};')
        dataloader = DataLoader(ds, batch_size=4, shuffle=True)
        for bidx, (X, y) in enumerate(dataloader):
            print(f'{bidx}: X: {X.shape}; y: {y.shape};')
            print(f'{X}')
            print(f'***************')
            print(f'{y}')
            if bidx > 5:
                break